"""
cms识别模块
"""

# 导入模块
from init import *


class WebCmsScan:
    def __init__(self, url, filepath=None, save_path=None):
        """
        定义参数
        """
        self.file_path = filepath   # 文件路径
        self.cms_dict = None    # 扫描信息
        self.url = url          # 扫描url
        self.current_url = None     # 多线程扫描的url
        self.request = None     # 扫描后返回的信息
        self.save_path = save_path  # 保存文件的路径
        self.result = None

    def _web_cms_scan(self, url):
        """
        调用接口，进行cms识别
        接口地址：http://whatweb.bugscaner.com/api.go
        接口地址：https://api.webscan.cc/?action=getip&domain=域名
        :param self:
        :param url: 扫描目标
        :return:
        """
        warnings.filterwarnings("ignore")  # 不输出警告信息
        # print(self.urls)
        # 伪造请求头
        #        扫描服务
        response = requests.get(url, verify=False)
        self.cms_dict = {"url": response.url, "text": response.text, "headers": dict(response.headers)}
        self.cms_dict = json.dumps(self.cms_dict)
        self.cms_dict = self.cms_dict.encode()
        self.cms_dict = zlib.compress(self.cms_dict)
        data = {"info": self.cms_dict}
        self.request = requests.post("http://whatweb.bugscaner.com/api.go", files=data, timeout=1.16)
        # 获取网站IP
        self.result = requests.get("https://api.webscan.cc/?action=getip&domain="+url)
        self._show_info()    # 输出信息

    def scan_options(self):
        """
        判断是读取文件还是直接扫描
        :return:
        """
        # print(self.file_path)
        if self.file_path:
            with open(self.file_path, 'r') as fp:    # 打开文件
                self.url = fp.read().split()         # 获取文件
                # print(self.url)
                for i in range(len(self.url)):       # 循环判断url有没有http://，没有就加上
                    if 'http://' not in self.url[i] or 'https://' not in self.url[i]:
                        self.url[i] = 'http://' + self.url[i]
                        # 调用写好的公共多线程参数
                        self._web_cms_scan(self.url)
        else:
            self._web_cms_scan(self.url)
            if self.save_path:
                self.save_file()

    def _show_info(self):
        """
        输出信息
        :return:
        """
        # 将扫描返回的json变成字典
        rs = dict(self.request.json())
        res = dict(self.result.json())
        warn_color.print('今日剩余的扫描次数：' + self.request.headers['X-RateLimit-Remaining'])
        if self.save_path:
            self.save_file()
        # 遍历字典输出信息
        if len(rs) > 1:
            success_color.print(f'[*]{rs.get("url")} 指纹信息如下：')
            success_color.print(f'[+]\tip地址：{res.get("ip")}')
            success_color.print(f'[+]\tip归属地：{res.get("info")}')
            for key in rs:
                if key != 'url':
                    success_color.print(f'[+]\t{key}：{rs[key]}')
        else:
            if rs.get("url"):
                error_color.print(f'[x]请求{rs.get("url")}失败')
            else:
                error_color.print(f'[x]请求{self.current_url}失败')

    def save_file(self):
        """
        保存并写入文件
        :return:
        """
        if self.save_path:
            with open(self.save_path, 'w', encoding='utf-8') as fp:
                fp.write(json.dumps(self.request.json(), indent=4, ensure_ascii=False))
            with open(self.save_path, 'a', encoding='utf-8') as fp:
                fp.write(json.dumps(self.result.json(), indent=4, ensure_ascii=False))
            success_color.print('保存成功\t', os.getcwd()+'\\'+self.save_path)
